import logging
from typing import List

import pandas as pd
import google.protobuf.any_pb2 as any_pb2

import coin.proto.coin_data_replay_pb2 as cdr_pb2
import coin2.service.account.account_info_service_pb2 as ais_pb2
from coin.base.timestamp import get_timestamp
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
from coin.base.datetime_util import to_datetime
from coin.support.data_replay.reader.interface import (
    TopicData,
    ITopicDataReader)
from coin.support.data_replay.reader.topic_data_reader import (
    _TopicDataTimestampContainer,
    DataMergeReader,
    TopicDataReader,)
from coin.support.data_replay.util.interval_container import IntervalContainer
from coin.support.pta.util.info_util import get_strat_info_list
from coin.util.queue.config import KafkaConfig
from xunkemgmt_client.client.util.query_util import (
    query_strategies)


def _validate_driver_config(driver_config: cdr_pb2.DataReplayDriverConfig):
  assert driver_config.HasField('start_timestamp')
  assert driver_config.HasField('end_timestamp')
  topic_set = set()
  for request in driver_config.requests:
    assert request.topic_type not in topic_set, request.topic_type
    topic_set.add(request.topic_type)


def _generate_strat_topic_data_requests(
    start_timestamp: int, end_timestamp: int, request: cdr_pb2.DataReplayRequest
) -> List[cdr_pb2.TopicDataRequest]:
  topic_data_requests = []
  strategies = query_strategies(
      strategy_ids=request.strategy_ids,
      strategy_names=request.strategy_names,
      strategy_groups=request.strategy_groups,
      traders=request.traders,
      business_units=request.business_units,
      as_proto=True)
  strategy_names = [strategy.strategy_name for strategy in strategies]
  start_date = to_datetime(start_timestamp).date()
  end_date = to_datetime(end_timestamp).date()
  if request.HasField('root_dir'):
    strat_infos = get_strat_info_list(
        start_date=start_date,
        end_date=end_date,
        strategy_list=strategy_names,
        machine_list=request.machines,
        strat_log_root_dir=request.root_dir)
    for strat_info in strat_infos:
      topic_data_request = cdr_pb2.TopicDataRequest(
          topic_type=request.topic_type,
          root_dir=request.root_dir,
          start_timestamp=start_timestamp,
          end_timestamp=end_timestamp,
          machine=strat_info.machine,
          strategy_name=strat_info.strategy_name
      )
      topic_data_requests.append(topic_data_request)
  if request.HasField('kafka_config_file'):
    strat_infos = get_strat_info_list(
        start_date=start_date,
        end_date=end_date,
        strategy_list=strategy_names,
        machine_list=request.machines,
        kafka_config=KafkaConfig.from_cmd_config(request.kafka_config_file))
    for strat_info in strat_infos:
      topic_data_request = cdr_pb2.TopicDataRequest(
          topic_type=request.topic_type,
          kafka_config_file=request.kafka_config_file,
          start_timestamp=start_timestamp,
          end_timestamp=end_timestamp,
          machine=strat_info.machine,
          strategy_name=strat_info.strategy_name
      )
      topic_data_requests.append(topic_data_request)
  return topic_data_requests


def _generate_topic_data_requests(
    start_timestamp: int, end_timestamp: int, request: cdr_pb2.DataReplayRequest):
  if request.topic_type in (cdr_pb2.STRATEGY, cdr_pb2.TELEMETRY):
    topic_data_requests = _generate_strat_topic_data_requests(
        start_timestamp, end_timestamp, request)
  elif request.topic_type in (cdr_pb2.BALANCE_REALTIME_HISTORIES,):
    topic_data_requests = to_list(
        cdr_pb2.TopicDataRequest(topic_type=request.topic_type,
                                 start_timestamp=start_timestamp,
                                 end_timestamp=end_timestamp,
                                 grpc_request=request.grpc_request,))
  else:
    raise ValueError('unknown topic type: %s' % request.topic_type)
  return topic_data_requests


def _generate_topic_data_readers(
    driver_config: cdr_pb2.DataReplayDriverConfig
) -> List[ITopicDataReader]:
  readers = []
  _validate_driver_config(driver_config)
  for request in driver_config.requests:
    topic_data_requests = _generate_topic_data_requests(
        driver_config.start_timestamp,
        driver_config.end_timestamp,
        request
    )
    for topic_data_request in topic_data_requests:
      readers.append(TopicDataReader(topic_data_request))
  return readers


class DataReplayHandler(object):
  def on_replay_start(self, ts: int):
    pass

  def on_replay_end(self, ts: int):
    pass

  def on_strat_log(self,
                   topic_info: cdr_pb2.TopicDataInfo,
                   ts: int,
                   data: StrategyLog):
    pass

  def on_og_log(self,
                topic_info: cdr_pb2.TopicDataInfo,
                ts: int,
                data: OrderGatewayLog):
    pass

  def on_interval_feed(self,
                       topic_info: cdr_pb2.TopicDataInfo,
                       ts: int,
                       data: pd.DataFrame):
    pass

  def on_balance_realtime_histories(self,
                                    topic_info: cdr_pb2.TopicDataInfo,
                                    ts: int,
                                    data: ais_pb2.QueryBalanceResponseProto):
    pass

  def on_timestamp(self, ts: int):
    pass


class DataReplayDriver(object):
  def __init__(self,
               driver_config: cdr_pb2.DataReplayDriverConfig,
               handler: DataReplayHandler):
    self._readers = _generate_topic_data_readers(driver_config)
    self._driver_config = driver_config
    self._handler = handler
    start_timestamp = driver_config.start_timestamp
    end_timestamp = driver_config.end_timestamp
    if driver_config.HasField('interval_sec'):
      frequency = driver_config.interval_sec * 10**9
    else:
      frequency = end_timestamp - start_timestamp
    self._interval_container = IntervalContainer(
        start_timestamp,
        end_timestamp,
        frequency)
    self._next_interval_ts = self._interval_container.get_interval(start_timestamp)[1]

  def run(self):
    merge_reader = DataMergeReader(self._readers, _TopicDataTimestampContainer)
    prev_ts = 0
    self._handler.on_replay_start(get_timestamp())
    while True:
      data = merge_reader.read()
      curr_ts = data.topic_data_value.timestamp if data is not None \
                else self._driver_config.end_timestamp
      self._on_timestamp_change(prev_ts, curr_ts)
      if data is None:
        break
      self._on_topic_data(data)
      prev_ts = curr_ts
    self._handler.on_timestamp(self._driver_config.end_timestamp)
    self._handler.on_replay_end(get_timestamp())

  def _on_timestamp_change(self, prev_ts: int, curr_ts: int):
    if prev_ts > curr_ts:
      #logging.info('%s vs %s' % (to_datetime(prev_ts), to_datetime(curr_ts)))
      return
    if curr_ts > self._next_interval_ts:
      if curr_ts < self._driver_config.end_timestamp:
        self._next_interval_ts = self._interval_container.get_interval(curr_ts)[1]
      intervals = self._interval_container.get_enclosing_interval(
          prev_ts, curr_ts)
      for each_ts in intervals:
        assert each_ts >= prev_ts, (each_ts, prev_ts)
        self._handler.on_timestamp(each_ts)

  def _on_topic_data(self, data: TopicData):
    topic_data_info = data.topic_data_info
    topic_type = topic_data_info.topic_type
    topic_data_value = data.topic_data_value
    if topic_type in (cdr_pb2.STRATEGY, cdr_pb2.TELEMETRY):
      pb = StrategyLog()
      pb.ParseFromString(topic_data_value.value)
      self._handler.on_strat_log(topic_data_info, topic_data_value.timestamp, pb)
    elif topic_type in (cdr_pb2.BALANCE_REALTIME_HISTORIES,):
      self._handler.on_balance_realtime_histories(
          topic_data_info, topic_data_value.timestamp, topic_data_value.value)
    else:
      raise ValueError('unknown topic type: %s' % topic_type)


if __name__ == '__main__':
  import datetime
  import json
  from coin.base.datetime_util import to_timestamp_int
  import google.protobuf.json_format as json_format
  from coin.base.param_util import to_list

  start_timestamp = to_timestamp_int(datetime.datetime(2023, 6, 21))
  end_timestamp = to_timestamp_int(datetime.datetime(2023, 6, 21, 1))
  requests = []
  requests.append(cdr_pb2.DataReplayRequest(
      topic_type=cdr_pb2.STRATEGY,
      root_dir='/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_slim_proto_log',
      kafka_config_file='../../coin_deploy/support_monitor/config/kafka_aws_config.json',
      strategy_groups=['simple_arb'],)
  )
  grpc_request = ais_pb2.QueryBalanceRealtimeHistoriesRequestProto(
      market_type='Futures',
      exchange='Binance',
      owner='xunke00_swap',
      start_timestamp=start_timestamp,
      end_timestamp=end_timestamp)
  any = any_pb2.Any()
  any.Pack(grpc_request)
  requests.append(cdr_pb2.DataReplayRequest(
      topic_type=cdr_pb2.BALANCE_REALTIME_HISTORIES,
      grpc_request=any,)
  )
  driver_config = cdr_pb2.DataReplayDriverConfig(
      start_timestamp=start_timestamp,
      end_timestamp=end_timestamp,
      requests=requests,
  )
  class SimpleHandler(DataReplayHandler):
    def __init__(self):
      self.datas = []

    def _on_topic_data(self, topic_info, ts, data):
      topic_info = json_format.MessageToDict(topic_info, preserving_proto_field_name=True)
      data = json_format.MessageToDict(data, preserving_proto_field_name=True)
      hr_ts = to_datetime(ts).strftime('%Y%m%d %H:%M:%S.%f')
      data['topic_info'] = topic_info
      data['record_timestamp'] = ts
      data['record_timestamp_human_readable'] = hr_ts
      self.datas.append(data)

    def on_timestamp(self, ts: int):
      print(to_datetime(ts))

    def on_strat_log(self,
                     topic_info: cdr_pb2.TopicDataInfo,
                     ts: int,
                     data: StrategyLog):
      self._on_topic_data(topic_info, ts, data)

    def on_balance_realtime_histories(self,
                                      topic_info: cdr_pb2.TopicDataInfo,
                                      ts: int,
                                      data: StrategyLog):
      self._on_topic_data(topic_info, ts, data)

  handler = SimpleHandler()
  driver = DataReplayDriver(driver_config, handler)
  driver.run()
  records = '[\n' + ',\n'.join(json.dumps(data) for data in handler.datas) + '\n]'
  print(records)
